import torch
import cv2
import matplotlib.pyplot as plt
from pathlib import Path

# 加载预训练模型
#model = torch.hub.load('/Users/michael/codes_ai/python/ai_demo_python/files/yolov5s.pt', 'yolov5s', pretrained=True)
model = torch.hub.load('ultralytics/yolov5', 'custom', path='/Users/michael/codes_ai/python/ai_demo_python/files/yolov5s.pt')
# 设置设备（自动选择GPU或CPU）
device = 'cuda' if torch.cuda.is_available() else 'mps'
model.to(device)

# 图像路径（替换为你自己的图片路径）
image_path = 'zidane.jpg'  # 示例图片
if not Path(image_path).exists():
    # 下载示例图片
    torch.hub.download_url_to_file(
        'https://ultralytics.com/images/zidane.jpg',
        image_path
    )

# 读取图像
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转换为RGB格式

# 进行推理
results = model(image_rgb)

# 解析结果
detections = results.pandas().xyxy[0]  # 转换为Pandas DataFrame

# 打印检测结果
print("检测到的对象:")
print(detections[['name', 'confidence', 'xmin', 'ymin', 'xmax', 'ymax']])

# 可视化结果
results.render()  # 在图像上绘制边界框和标签
output_image = results.ims[0]  # 获取带标注的图像

# 显示结果
plt.figure(figsize=(12, 8))
plt.imshow(output_image)
plt.axis('off')
plt.title('YOLOv5 目标检测结果')
plt.show()

# 保存结果
output_path = 'detection_result.jpg'
cv2.imwrite(output_path, cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR))
print(f"结果已保存至: {output_path}")